import numpy as np
import torch.nn.functional as F

from .unet import UNet
from .ddpm import GaussianDiffusion


def generate_linear_schedule(T, low, high):
    return np.linspace(low, high, T)


def generate_cosine_schedule(T, s=0.008):
    def f(t, T):
        return (np.cos((t / T + s) / (1 + s) * np.pi / 2)) ** 2

    alphas = []
    f0 = f(0, T)

    for t in range(T + 1):
        alphas.append(f(t, T) / f0)

    betas = []

    for t in range(1, T + 1):
        betas.append(min(1 - alphas[t] / alphas[t - 1], 0.999))

    return np.array(betas)


def build_model(args):

    if args.model == "ddpm":
        #  定义激活函数
        activations = {
            "relu": F.relu,
            "mish": F.mish,
            "silu": F.silu,
        }

        model = UNet(
            img_channels=args.in_channels,
            base_channels=args.base_channels,
            channel_mults=args.channel_mults,
            time_emb_dim=args.time_emb_dim,
            norm=args.norm,
            dropout=args.dropout,
            activation=activations[args.activation],
            attention_resolutions=args.attention_resolutions,
            num_classes=None if not args.use_labels else 10,
            initial_pad=0,
        )

        if args.schedule == "cosine":
            betas = generate_cosine_schedule(args.num_timesteps)
        else:
            betas = generate_linear_schedule(
                args.num_timesteps,
                args.schedule_low * 1000 / args.num_timesteps,
                args.schedule_high * 1000 / args.num_timesteps,
            )

        model = GaussianDiffusion(
            model=model, 
            img_size=(32, 32), 
            img_channels=args.in_channels, 
            num_classes=10,
            betas=betas,
            ema_decay=args.ema_decay,
            ema_update_rate=args.ema_update_rate,
            ema_start=2000,
            loss_type=args.loss_type,
        )

    return model





